package runner

import (
    "fmt"
    "net/http"
    "regexp"
    "strings"
    "sync"
    "time"

    "github.com/zan8in/afrog/v3/pkg/config"
    "github.com/zan8in/afrog/v3/pkg/protocols/http/retryhttpclient"
    "github.com/zan8in/afrog/v3/pkg/result"

    "github.com/google/cel-go/checker/decls"
    "github.com/zan8in/afrog/v3/pkg/poc"
    "github.com/zan8in/afrog/v3/pkg/proto"
    "github.com/zan8in/afrog/v3/pkg/utils"
    "gopkg.in/yaml.v2"
)

var MMutex = &sync.Mutex{}

type Checker struct {
	Options *config.Options
	// OriginalRequest *http.Request
	VariableMap map[string]any
	Result      *result.Result
	CustomLib   *CustomLib
}

func (c *Checker) Check(target string, pocItem *poc.Poc) (err error) {
	defer func() {
		if r := recover(); r != nil {
			c.Result.IsVul = false
		}
	}()

	if pocItem.IsHTTPType() {
		if target, err = c.checkURL(target); err != nil {
			c.Result.IsVul = false
			return err
		}
	}

	if !strings.HasPrefix(target, "http://") && !strings.HasPrefix(target, "https://") {
		target = "https://" + target
	}

	c.Result.Target = target
	c.Result.PocInfo = pocItem

	matchCondition := ""
	if strings.Contains(pocItem.Expression, "&&") && !strings.Contains(pocItem.Expression, "||") {
		matchCondition = poc.STOP_IF_FIRST_MISMATCH
	}
	if strings.Contains(pocItem.Expression, "||") && !strings.Contains(pocItem.Expression, "&&") {
		matchCondition = poc.STOP_IF_FIRST_MATCH
	}

	originReq, err := http.NewRequest("GET", target, nil)
	if err != nil {
		c.Result.IsVul = false
		return err
	}
	tempRequest, err := retryhttpclient.ParseRequest(originReq)
	if err != nil {
		c.Result.IsVul = false
		return err
	}
	c.VariableMap["request"] = tempRequest

	if len(pocItem.Set) > 0 {
		c.UpdateVariableMap(pocItem.Set)
	}

	if len(pocItem.Payloads.Payloads) > 0 {
		c.UpdateVariableMap(pocItem.Payloads.Payloads)
	}

	for _, ruleMap := range pocItem.Rules {
		k := ruleMap.Key
		rule := ruleMap.Value

		if rule.BeforeSleep != 0 {
			time.Sleep(time.Duration(rule.BeforeSleep) * time.Second)
		}

		// 预处理：让 rules 的 path/headers/body/host/raw/data 支持 {{...}} CEL 运算
		c.preRenderRuleRequest(&rule.Request)

        isMatch := false
        reqType := strings.ToLower(rule.Request.Type)

        if len(rule.Request.Raw) > 0 {
            err = RawHTTPExecutor{}.Execute(target, rule, c.Options, c.VariableMap)
        } else {
            exec, ok := executors[reqType]
            if !ok {
                exec = HTTPExecutor{}
            }
            err = exec.Execute(target, rule, c.Options, c.VariableMap)
        }

		if err == nil {
			if len(rule.Expressions) > 0 {
				// multiple expressions
				for _, expression := range rule.Expressions {
					evalResult, err := c.CustomLib.RunEval(expression, c.VariableMap)
					if err == nil {
						isMatch = evalResult.Value().(bool)
						if isMatch {
							if name := checkExpression(expression); len(name) > 0 {
								pocItem.Id = name
								pocItem.Info.Name = name
							}
							break
						}
					}
				}
			} else {
				// single expression
				evalResult, err := c.CustomLib.RunEval(rule.Expression, c.VariableMap)
				if err == nil {
					isMatch = evalResult.Value().(bool)
				}
			}
		}

		c.CustomLib.WriteRuleFunctionsROptions(k, isMatch)

		if len(rule.Output) > 0 && isMatch {
			c.UpdateVariableMap(rule.Output)
		}

		if len(rule.Extractors) > 0 && isMatch {
			c.UpdateVariableMapExtractor(rule.Extractors)
		}

		pocRstTemp := result.PocResult{IsVul: isMatch}
		if c.VariableMap["response"] != nil {
			pocRstTemp.ResultResponse = c.VariableMap["response"].(*proto.Response)
		}
		if c.VariableMap["request"] != nil {
			pocRstTemp.ResultRequest = c.VariableMap["request"].(*proto.Request)
		}
		if c.VariableMap["fulltarget"] != nil {
			pocRstTemp.FullTarget = c.VariableMap["fulltarget"].(string)
			c.Result.FullTarget = c.VariableMap["fulltarget"].(string)
		}
		if c.VariableMap["target"] != nil {
			c.Result.Target = c.VariableMap["target"].(string)
		}
		c.Result.AllPocResult = append(c.Result.AllPocResult, &pocRstTemp)

		if rule.StopIfMismatch && !isMatch {
			c.Result.IsVul = false
			return err
		}

		if rule.StopIfMatch && isMatch {
			if len(pocItem.Extractors) > 0 {
				c.UpdateVariableMapExtractor(pocItem.Extractors)
			}
			c.Result.IsVul = true
			return err
		}

		if matchCondition == poc.STOP_IF_FIRST_MISMATCH && !isMatch {
			c.Result.IsVul = false
			return err
		}

		if matchCondition == poc.STOP_IF_FIRST_MATCH && isMatch {
			if len(pocItem.Extractors) > 0 {
				c.UpdateVariableMapExtractor(pocItem.Extractors)
			}
			c.Result.IsVul = true
			return err
		}

	}

	if len(pocItem.Extractors) > 0 {
		c.UpdateVariableMapExtractor(pocItem.Extractors)
	}

	isVul, err := c.CustomLib.RunEval(pocItem.Expression, c.VariableMap)
	if err != nil {
		c.Result.IsVul = false
		return err
	}

	c.Result.IsVul = isVul.Value().(bool)

	return err
}

func (c *Checker) checkURL(target string) (string, error) {

	// if target check num more than MaxCheckNum
	tcount := c.Options.Targets.Num(target)

	if tcount == ActiveTarget {
		return target, nil
	}

	if tcount > c.Options.MaxHostError {
		return "", fmt.Errorf("%s is blacklisted", target)
	}

	// if target is not url, then check again
	if !utils.IsURL(target) {
		if newtarget, err := retryhttpclient.CheckProtocol(target); err == nil {
			if k := c.Options.Targets.Key(target); k >= 0 {
				c.Options.Targets.Update(k, newtarget)
				c.Options.Targets.SetNum(newtarget, ActiveTarget)
			}
			return newtarget, nil
		}

		c.Options.Targets.UpdateNum(target, 1)
		return target, fmt.Errorf("%s check protocol falied", target)
	}

	// if target is url more than zero, then check protocol against
	if c.Options.Targets.Num(target) >= 0 {
		if newtarget, err := retryhttpclient.CheckProtocol(target); err == nil {
			c.Options.Targets.SetNum(newtarget, ActiveTarget)
			return newtarget, nil
		}

		c.Options.Targets.UpdateNum(target, 1)
		return target, fmt.Errorf("%s no response", target)
	}

	return target, nil
}

func (c *Checker) UpdateVariableMap(args yaml.MapSlice) {
	for _, item := range args {
		key := item.Key.(string)

		// 新增：根据 YAML 值的实际类型分支处理
		switch v := item.Value.(type) {
		case string:
			// oob() 函数特殊处理
			if v == "oob()" {
				c.VariableMap[key] = c.oob()
				c.CustomLib.UpdateCompileOption(key, decls.NewObjectType("proto.OOB"))
				continue
			}

			// 原有字符串路径：走 CEL 求值
			out, err := c.CustomLib.RunEval(v, c.VariableMap)
			if err != nil {
				// fixed set string failed bug
				c.VariableMap[key] = fmt.Sprintf("%v", v)
				c.CustomLib.UpdateCompileOption(key, decls.String)
				continue
			}
			switch value := out.Value().(type) {
			case *proto.UrlType:
				c.VariableMap[key] = utils.UrlTypeToString(value)
				c.CustomLib.UpdateCompileOption(key, decls.NewObjectType("proto.UrlType"))
			case int64:
				c.VariableMap[key] = int(value)
				c.CustomLib.UpdateCompileOption(key, decls.Int)
			case map[string]string:
				c.VariableMap[key] = value
				c.CustomLib.UpdateCompileOption(key, StrStrMapType)
			default:
				c.VariableMap[key] = fmt.Sprintf("%v", out)
				c.CustomLib.UpdateCompileOption(key, decls.String)
			}

		case int:
			c.VariableMap[key] = v
			c.CustomLib.UpdateCompileOption(key, decls.Int)
			continue
		case int64:
			c.VariableMap[key] = int(v)
			c.CustomLib.UpdateCompileOption(key, decls.Int)
			continue
		case float64:
			c.VariableMap[key] = v
			c.CustomLib.UpdateCompileOption(key, decls.Double)
			continue
		case bool:
			c.VariableMap[key] = v
			c.CustomLib.UpdateCompileOption(key, decls.Bool)
			continue
		default:
			// 其他类型统一按字符串存（保证兼容）
			c.VariableMap[key] = fmt.Sprintf("%v", v)
			c.CustomLib.UpdateCompileOption(key, decls.String)
			continue
		}
	}
}

func (c *Checker) UpdateVariableMapExtractor(extractors []poc.Extractors) {
	for _, v := range extractors {
		tpe := v.Type
		extMap := v.Extractor
		if len(extMap) == 0 {
			continue
		}
		for _, item := range extMap {
			key := item.Key.(string)
			value := item.Value.(string)

			if tpe == "word" {
				new := setVariableMap(value, c.VariableMap)
				c.VariableMap[key] = new
				c.CustomLib.UpdateCompileOption(key, decls.String)
				c.Result.Extractor = append(c.Result.Extractor, yaml.MapItem{Key: key, Value: new})
				continue
			}

			out, err := c.CustomLib.RunEval(value, c.VariableMap)
			if err != nil {
				continue
			}

			switch value := out.Value().(type) {
			case map[string]string:
				c.VariableMap[key] = value
				c.CustomLib.UpdateCompileOption(key, StrStrMapType)
				c.Result.Extractor = append(c.Result.Extractor, yaml.MapItem{Key: key, Value: value})
			case string:
				c.VariableMap[key] = fmt.Sprintf("%v", out)
				c.CustomLib.UpdateCompileOption(key, decls.String)
				c.Result.Extractor = append(c.Result.Extractor, yaml.MapItem{Key: key, Value: value})
			}

		}
	}

}

func (c *Checker) oob() *proto.OOB {
	if OOB == nil {
		return &proto.OOB{}
	}

	vdomains := OOB.GetValidationDomain()

	return &proto.OOB{
		Filter:       vdomains.Filter,
		HTTP:         vdomains.HTTP,
		DNS:          vdomains.DNS,
		ProtocolHTTP: "http",
		ProtocolDNS:  "dns",
	}
}

// func (c *Checker) newRerverse() *proto.Reverse {

// 	urlStr := ""
// 	// sub := utils.CreateRandomString(20)

// 	// 使用反连平台优先权逻辑如下：
// 	// 自建eye反连平台 > ceye反连平台 > eyes.sh反连平台
// 	// @edit 2021.11.29 21:50
// 	// 关联代码 celprogram.go line-596
// 	// if config.ReverseEyeShLive && config.ReverseEyeHost != "eyes.sh" {
// 	// 	urlStr = fmt.Sprintf("http://%s.%s", sub, config.ReverseEyeDomain)
// 	// } else if config.ReverseCeyeLive {
// 	// 	urlStr = fmt.Sprintf("http://%s.%s", sub, config.ReverseCeyeDomain)
// 	// } else if config.ReverseEyeShLive {
// 	// 	urlStr = fmt.Sprintf("http://%s.%s", sub, config.ReverseEyeDomain)
// 	// }

// 	u, _ := url.Parse(urlStr)
// 	return &proto.Reverse{
// 		Url:                utils.ParseUrl(u),
// 		Domain:             u.Hostname(),
// 		Ip:                 u.Host,
// 		IsDomainNameServer: false,
// 	}
// }

// func (c *Checker) newJNDI() *proto.Reverse {
// 	// randomstr := utils.CreateRandomString(22)
// 	// urlStr := fmt.Sprintf("http://%s:%s/%s", config.ReverseJndi, config.ReverseLdapPort, randomstr)
// 	// u, _ := url.Parse(urlStr)
// 	// url := utils.ParseUrl(u)
// 	// return &proto.Reverse{
// 	// 	Url:                url,
// 	// 	Domain:             u.Hostname(),
// 	// 	Ip:                 config.ReverseJndi,
// 	// 	IsDomainNameServer: false,
// 	// }
// 	return &proto.Reverse{}
// }

func checkExpression(expression string) string {
	if strings.Contains(expression, "!= \"\"") {
		pos := strings.Index(expression, "!= \"\"")
		name := strings.Trim(strings.TrimSpace(expression[:pos]), "\"")
		return name
	}
	return ""

}

func setVariableMap(find string, variableMap map[string]any) string {
	for k, v := range variableMap {
		_, isMap := v.(map[string]string)
		if isMap {
			continue
		}
		newstr := fmt.Sprintf("%v", v)
		oldstr := "{{" + k + "}}"
		if !strings.Contains(find, oldstr) {
			continue
		}
		find = strings.ReplaceAll(find, oldstr, newstr)
	}
	return find
}

func (c *Checker) renderCELPlaceholders(s string) string {
	// 为空直接返回
	if len(s) == 0 {
		return s
	}
	re := regexp.MustCompile(`\{\{(.+?)\}\}`)
	return re.ReplaceAllStringFunc(s, func(m string) string {
		// 提取 {{ ... }} 内的表达式文本
		expr := strings.TrimSpace(m[2 : len(m)-2])
		// 优先走 CEL 求值
		out, err := c.CustomLib.RunEval(expr, c.VariableMap)
		if err != nil {
			// 求值失败，保留原占位符，后续仍可由 setVariableMap 做简单变量替换
			return m
		}
		switch v := out.Value().(type) {
		case *proto.UrlType:
			return utils.UrlTypeToString(v)
		case []byte:
			return string(v)
		default:
			return fmt.Sprintf("%v", v)
		}
	})
}

// 对当前 rule 的请求字段进行预渲染（仅替换能成功求值的表达式）
func (c *Checker) preRenderRuleRequest(req *poc.RuleRequest) {
    // HTTP(S)/RAW/NETX 通用字段
    req.Path = c.renderCELPlaceholders(strings.TrimSpace(req.Path))
    req.Host = c.renderCELPlaceholders(strings.TrimSpace(req.Host))
    req.Body = c.renderCELPlaceholders(strings.TrimSpace(req.Body))
    req.Raw = c.renderCELPlaceholders(strings.TrimSpace(req.Raw))
    req.Data = c.renderCELPlaceholders(strings.TrimSpace(req.Data))

    // CEL 渲染后做一次简单 {{var}} 替换作为兜底
    req.Path = setVariableMap(req.Path, c.VariableMap)
    req.Host = setVariableMap(req.Host, c.VariableMap)
    req.Body = setVariableMap(req.Body, c.VariableMap)
    req.Raw = setVariableMap(req.Raw, c.VariableMap)
    req.Data = setVariableMap(req.Data, c.VariableMap)

    // headers 逐项处理（深拷贝避免并发写入共享 map）
    if req.Headers != nil {
        newHeaders := make(map[string]string, len(req.Headers))
        for hk, hv := range req.Headers {
            h := c.renderCELPlaceholders(strings.TrimSpace(hv))
            newHeaders[hk] = setVariableMap(h, c.VariableMap)
        }
        req.Headers = newHeaders
    }
}
